# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2025/1/27
# @Function: QA数据库查询工具
import os
import sys

import pandas as pd

from test.data_import.qa_query_service import exec_nlp_query

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from loguru import logger

if __name__ == "__main__":
    logger.debug("QA查询工具测试")
    logger.debug("=" * 50)

    test_questions = [
        "大伙房水库是几级水库"  # 生态关注
    ]
    result_list = []
    for question in test_questions:
        markdown_output = exec_nlp_query(question, '21100150')
        result_list.append(
            {
                "question": question,
                "answser": markdown_output
            }
        )
    pd.DataFrame(result_list).to_excel("test_result.xlsx", index=False)
